/**
 * Copyright (c) 2011-2014, hubin (jobob@qq.com).
 * <p>
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package city.spring.configure.mybatis;

import city.spring.modules.log.entity.SqlLogEntity;
import city.spring.modules.log.mq.SqlLogNoticeTopic;
import city.spring.utils.RequestUtils;
import city.spring.utils.UserUtils;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.handlers.AbstractSqlParserHandler;
import com.google.common.collect.EvictingQueue;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.logging.Log;
import org.apache.ibatis.logging.LogFactory;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.ResultHandler;
import org.springframework.integration.support.MessageBuilder;
import org.springframework.messaging.Message;
import org.springframework.stereotype.Component;
import org.springframework.util.ReflectionUtils;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.security.Principal;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * <p>
 * 拦截系统的SQL执行记录，主要记录系统对数据库的增删查改操作，以便后期追踪数据更改问题
 * </p>
 *
 * @author HouKunLin
 * @date 2020-08-5 13:56:34
 * @see <a href="https://github.com/baomidou/mybatis-plus/blob/2.x/mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/plugins/PerformanceInterceptor.java">PerformanceInterceptor.java</a>
 */
@Intercepts({
        @Signature(type = StatementHandler.class, method = "update", args = {Statement.class}),
        @Signature(type = StatementHandler.class, method = "batch", args = {Statement.class}),
        @Signature(type = StatementHandler.class, method = "query", args = {Statement.class, ResultHandler.class}),
})
@Component
public class ExecuteSqlInterceptor extends AbstractSqlParserHandler implements Interceptor {
    private static final Log logger = LogFactory.getLog(ExecuteSqlInterceptor.class);
    private static final String DruidPooledPreparedStatement = "com.alibaba.druid.pool.DruidPooledPreparedStatement";
    private static final String T4CPreparedStatement = "oracle.jdbc.driver.T4CPreparedStatement";
    private static final String OraclePreparedStatementWrapper = "oracle.jdbc.driver.OraclePreparedStatementWrapper";
    private static final String P6spyPreparedStatementWrapper = "com.p6spy.engine.wrapper.PreparedStatementWrapper";
    private final SqlLogNoticeTopic sqlLogNoticeTopic;
    /**
     * 只记录以下几种类型的SQL执行记录
     */
    private final List<SqlCommandType> commandTypes = new ArrayList<SqlCommandType>() {{
        add(SqlCommandType.INSERT);
        add(SqlCommandType.UPDATE);
        add(SqlCommandType.DELETE);
    }};
    private final EvictingQueue<Message<SqlLogEntity>> queue = EvictingQueue.create(10);
    private Method oracleGetOriginalSqlMethod;
    private Method druidGetSqlMethod;
    private Field p6spyGetSqlField;

    public ExecuteSqlInterceptor(SqlLogNoticeTopic sqlLogNoticeTopic) {
        this.sqlLogNoticeTopic = sqlLogNoticeTopic;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        long start = System.currentTimeMillis();
        try {
            return invocation.proceed();
        } finally {
            SqlLogEntity entity = new SqlLogEntity();
            entity.setDuration(System.currentTimeMillis() - start);
            HttpServletRequest request = RequestUtils.getRequest();
            if (request != null) {
                entity.setIp(RequestUtils.getRequestIp(request));
                entity.setRequestCode(String.valueOf(request.hashCode()));
            }

            Principal principal = UserUtils.getPrincipal();
            if (principal != null) {
                entity.setUserId(principal.getName());
            }

            Object[] args = invocation.getArgs();
            Statement statement = getStatement(args[0]);
            entity.setSqlText(obtainSql(getStatementSql(statement)));

            StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
            MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
            MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
            if (mappedStatement != null) {
                entity.setType(mappedStatement.getSqlCommandType().name());
            }

            Message<SqlLogEntity> message = MessageBuilder.withPayload(entity).build();
            try {
                sqlLogNoticeTopic.output().send(message);
                if (!queue.isEmpty()) {
                    Message<SqlLogEntity> remove = queue.remove();
                    if (remove != null) {
                        sqlLogNoticeTopic.output().send(remove);
                    }
                }
            } catch (Exception ignore) {
                queue.add(message);
            }
        }
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    public Method getMethodRegular(Class<?> clazz, String methodName) {
        if (Object.class.equals(clazz)) {
            return null;
        }
        for (Method method : clazz.getDeclaredMethods()) {
            if (method.getName().equals(methodName)) {
                return method;
            }
        }
        return getMethodRegular(clazz.getSuperclass(), methodName);
    }

    /**
     * 获取sql语句开头部分
     *
     * @param sql SQL语句
     * @return int
     */
    private int indexOfSqlStart(String sql) {
        String upperCaseSql = sql.toUpperCase();
        Set<Integer> set = new HashSet<>();
        set.add(upperCaseSql.indexOf("SELECT "));
        set.add(upperCaseSql.indexOf("UPDATE "));
        set.add(upperCaseSql.indexOf("INSERT "));
        set.add(upperCaseSql.indexOf("DELETE "));
        set.remove(-1);
        if (CollectionUtils.isEmpty(set)) {
            return -1;
        }
        List<Integer> list = new ArrayList<>(set);
        list.sort(Integer::compareTo);
        return list.get(0);
    }

    /**
     * 获得当前的 Statement 对象
     *
     * @param firstArg 拦截器的参数
     * @return Statement
     * @author 侯坤林
     * @date 2020-08-5 13:52:33
     */
    public Statement getStatement(Object firstArg) {
        Statement statement;
        if (Proxy.isProxyClass(firstArg.getClass())) {
            statement = (Statement) SystemMetaObject.forObject(firstArg).getValue("h.statement");
        } else {
            statement = (Statement) firstArg;
        }
        MetaObject stmtMetaObj = SystemMetaObject.forObject(statement);
        try {
            statement = (Statement) stmtMetaObj.getValue("stmt.statement");
        } catch (Exception ignored) {
        }
        //Hikari
        if (stmtMetaObj.hasGetter("delegate")) {
            try {
                statement = (Statement) stmtMetaObj.getValue("delegate");
            } catch (Exception ignored) {
            }
        }
        return statement;
    }

    /**
     * 获得当前正在执行的完整SQL语句
     *
     * @param statement statement对象
     * @return SQL语句（带有前缀信息的）
     * @author 侯坤林
     * @date 2020-08-5 13:52:33
     */
    private String getStatementSql(Statement statement) {
        String stmtClassName = statement.getClass().getName();
        if (P6spyPreparedStatementWrapper.equals(stmtClassName)) {
            try {
                if (p6spyGetSqlField == null) {
                    p6spyGetSqlField = statement.getClass().getDeclaredField("delegate");
                    p6spyGetSqlField.setAccessible(true);
                }
                Object value = ReflectionUtils.getField(p6spyGetSqlField, statement);
                if (value instanceof Statement) {
                    return getStatementSql((Statement) value);
                }
            } catch (NoSuchFieldException ignored) {
            }
        }
        if (DruidPooledPreparedStatement.equals(stmtClassName)) {
            try {
                if (druidGetSqlMethod == null) {
                    Class<?> clazz = Class.forName(DruidPooledPreparedStatement);
                    druidGetSqlMethod = clazz.getMethod("getSql");
                }
                Object stmtSql = druidGetSqlMethod.invoke(statement);
                if (stmtSql instanceof String) {
                    return (String) stmtSql;
                }
            } catch (Exception ignored) {
            }
        } else if (T4CPreparedStatement.equals(stmtClassName)
                || OraclePreparedStatementWrapper.equals(stmtClassName)) {
            try {
                if (oracleGetOriginalSqlMethod != null) {
                    Object stmtSql = oracleGetOriginalSqlMethod.invoke(statement);
                    if (stmtSql instanceof String) {
                        return (String) stmtSql;
                    }
                } else {
                    Class<?> clazz = Class.forName(stmtClassName);
                    oracleGetOriginalSqlMethod = getMethodRegular(clazz, "getOriginalSql");
                    if (oracleGetOriginalSqlMethod != null) {
                        //OraclePreparedStatementWrapper is not a public class, need set this.
                        oracleGetOriginalSqlMethod.setAccessible(true);
                        if (oracleGetOriginalSqlMethod != null) {
                            Object stmtSql = oracleGetOriginalSqlMethod.invoke(statement);
                            if (stmtSql instanceof String) {
                                return (String) stmtSql;
                            }
                        }
                    }
                }
            } catch (Exception ignored) {
            }
        }
        return statement.toString();
    }

    /**
     * 获取能够直接执行的SQL语句（删除前缀信息）
     *
     * @param statementSql SQL语句（带有前缀信息的）
     * @return SQL语句（能够直接执行的SQL语句）
     * @author 侯坤林
     * @date 2020-08-5 13:52:33
     */
    private String obtainSql(String statementSql) {
        String originalSql = statementSql.replaceAll("[\\s]+", " ");
        int index = indexOfSqlStart(originalSql);
        if (index > 0) {
            originalSql = originalSql.substring(index);
        }
        return originalSql;
    }
}
